package com.project.fortuneteller.utils;

import cn.smartjavaai.common.config.Config;
import cn.smartjavaai.common.entity.R;
import cn.smartjavaai.common.entity.face.FaceSearchResult;
import cn.smartjavaai.common.enums.DeviceEnum;
import cn.smartjavaai.face.config.FaceDetConfig;
import cn.smartjavaai.face.config.FaceRecConfig;
import cn.smartjavaai.face.constant.FaceDetectConstant;
import cn.smartjavaai.face.entity.FaceRegisterInfo;
import cn.smartjavaai.face.entity.FaceSearchParams;
import cn.smartjavaai.face.enums.FaceDetModelEnum;
import cn.smartjavaai.face.enums.FaceRecModelEnum;
import cn.smartjavaai.face.enums.SimilarityType;
import cn.smartjavaai.face.factory.FaceDetModelFactory;
import cn.smartjavaai.face.factory.FaceRecModelFactory;
import cn.smartjavaai.face.model.facedect.FaceDetModel;
import cn.smartjavaai.face.model.facerec.FaceRecModel;
import cn.smartjavaai.face.vector.config.SQLiteConfig;
import cn.smartjavaai.face.vector.entity.FaceVector;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.node.ObjectNode;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.awt.image.BufferedImage;
import java.util.Date;
import java.util.List;

/**
 * 人脸识别模型
 * 文档地址：http://doc.smartjavaai.cn/
 * @author dwj
 */
public class FaceRecUtils {

    private static final Logger log = LoggerFactory.getLogger(FaceRecUtils.class);

    private static final ObjectMapper objectMapper = new ObjectMapper();

    // 设备类型
    public static final DeviceEnum device = DeviceEnum.CPU;

    public static final String MODAL_PATH = "D:\\www-environment\\smart_java_ai\\";

    static {
        //修改缓存路径
        Config.setCachePath(MODAL_PATH);
    }

    public static void init() {
        getFaceRecModelWithSQLiteConfig();
    }

    /**
     * 人脸更新
     */
    public static void updateFace(BufferedImage image, String userid, String role, String faceId) {
        try {
            //高精度模型，速度慢, 追求速度请更换高速模型
            FaceRecModel faceRecModel = getFaceRecModelWithSQLiteConfig();
            //等待加载人脸库结束
            while (!faceRecModel.isLoadFaceCompleted()){
                Thread.sleep(100);
            }
            //设置人脸注册的自定义元数据，本例中使用 JSON 格式存储用户信息
            ObjectNode metadataJson = objectMapper.createObjectNode();
            metadataJson.put("userid", userid);
            metadataJson.put("role", role);

            log.info("====================人脸更新==========================");
            FaceRegisterInfo updateInfo = new FaceRegisterInfo();
            //设置人脸注册的自定义元数据，本例中使用 JSON 格式存储用户信息
            updateInfo.setMetadata(objectMapper.writeValueAsString(metadataJson));
            //更新必须设置ID
            updateInfo.setId(faceId);
            faceRecModel.upsertFace(updateInfo, image);
            log.info("更新人脸成功");
        } catch (Exception e) {
            log.error("人脸特征比对异常: {}", e.getMessage(), e);
        }
    }
    /**
     * 人脸查询
     */
    public static FaceSearchResult queryFace(BufferedImage image) {
        try {
            //高精度模型，速度慢, 追求速度请更换高速模型
            FaceRecModel faceRecModel = getFaceRecModelWithSQLiteConfig();
            //等待加载人脸库结束
            while (!faceRecModel.isLoadFaceCompleted()){
                Thread.sleep(100);
            }

            log.info("====================人脸查询==========================");
            //特征提取（提取分数最高人脸特征）,适用于单人脸场景
            R<float[]> featureResult2 = faceRecModel.extractTopFaceFeature(image);
            if (featureResult2.isSuccess()) {
                log.info("人脸特征提取成功, {}", objectMapper.writeValueAsString(featureResult2.getData()));
            } else {
                log.info("人脸特征提取失败：{}", featureResult2.getMessage());
                return null;
            }
            FaceSearchParams faceSearchParams = new FaceSearchParams();
            faceSearchParams.setTopK(1);
            faceSearchParams.setThreshold(0.8f);
            List<FaceSearchResult> faceSearchResults = faceRecModel.search(featureResult2.getData(), faceSearchParams);
            log.info("人脸查询结果：{}", objectMapper.writeValueAsString(faceSearchResults));
            // 返回第一个人脸
            if (!faceSearchResults.isEmpty()) {
                return faceSearchResults.get(0);
            }
            return new FaceSearchResult(null, 0.0f, null);
        } catch (Exception e) {
            log.error("人脸查询异常: {}", e.getMessage(), e);
        }
        return null;
    }


    /**
     * 人脸删除
     */
    public static void deleteFace(String faceId) {
        try {
            //高精度模型，速度慢, 追求速度请更换高速模型
            FaceRecModel faceRecModel = getFaceRecModelWithSQLiteConfig();
            //等待加载人脸库结束
            while (!faceRecModel.isLoadFaceCompleted()){
                Thread.sleep(100);
            }
            log.info("====================人脸删除==========================");
            try {
                faceRecModel.removeRegister(faceId);
                log.info("人脸删除成功");
            } catch (Exception e) {
                log.info("人脸删除失败,可能不存在：{}", e.getMessage());
            }
        } catch (Exception e) {
            log.error("人脸删除异常: {}", e.getMessage(), e);
        }
    }

    /**
     * 注册人脸信息，传入人脸图片和用户信息
     * @param image
     * @param userid
     * @param role
     */
    public static String registerFace(BufferedImage image, String userid, String role) {
        try {
            //高精度模型，速度慢, 追求速度请更换高速模型
            FaceRecModel faceRecModel = getFaceRecModelWithSQLiteConfig();
            //等待加载人脸库结束
            while (!faceRecModel.isLoadFaceCompleted()){
                Thread.sleep(100);
            }
            log.info("====================人脸注册==========================");
            //特征提取（提取分数最高人脸特征）,适用于单人脸场景
            R<float[]> featureResult = faceRecModel.extractTopFaceFeature(image);
            if (featureResult.isSuccess()) {
                log.info("人脸特征提取成功: {}", objectMapper.writeValueAsString(featureResult.getData()));
            } else {
                log.info("人脸特征提取失败：{}", featureResult.getMessage());
                return null;
            }
            //人脸注册信息
            FaceRegisterInfo faceRegisterInfo = new FaceRegisterInfo();
            //设置人脸注册的自定义元数据，本例中使用 JSON 格式存储用户信息
            ObjectNode metadataJson = objectMapper.createObjectNode();
            metadataJson.put("userid", userid);
            metadataJson.put("role", role);
            faceRegisterInfo.setMetadata(objectMapper.writeValueAsString(metadataJson));
            //可自定义 ID，若未设置则自动生成。
//            faceRegisterInfo.setId("00001");
            // 查询是否人脸信息是否已存在

            //人脸注册，返回人脸库ID
            R<String> registerResult = faceRecModel.register(faceRegisterInfo, featureResult.getData());
            if (registerResult.isSuccess()) {
                log.info("注册成功：ID-{}", registerResult.getData());
            } else {
                log.info("注册失败：{}", registerResult.getMessage());
                return null;
            }

            return registerResult.getData();
        } catch (Exception e) {
            log.error("人脸注册异常: {}", e.getMessage(), e);
        }
        return null;
    }

    private static FaceDetModel faceDetModel = null;
    /**
     * 1 获取人脸检测模型（高精度模型）
     * 注意事项：
     * 1、高精度模型，识别准确度高，速度慢
     * 2、具体其他模型参数可以查看文档：http://doc.smartjavaai.cn/face.html
     * @return
     */
    public static FaceDetModel getProFaceDetModel() {
        if (faceDetModel != null) {
            return faceDetModel;
        }
        FaceDetConfig config = new FaceDetConfig();
        //人脸检测模型，SmartJavaAI提供了多种模型选择(更多模型，请查看文档)，切换模型需要同时修改modelEnum及modelPath
        config.setModelEnum(FaceDetModelEnum.RETINA_FACE);
        //下载模型并替换本地路径，下载地址：https://pan.baidu.com/s/10l22x5fRz_gwLr8EAHa1Jg?pwd=1234 提取码: 1234
        config.setModelPath(MODAL_PATH + "retinaface.pt");
        //只返回相似度大于该值的人脸,需要根据实际情况调整，分值越大越严格容易漏检，分值越小越宽松容易误识别
        config.setConfidenceThreshold(0.5f);
        //用于去除重复的人脸框，当两个框的重叠度超过该值时，只保留一个
        config.setNmsThresh(FaceDetectConstant.NMS_THRESHOLD);
        faceDetModel = FaceDetModelFactory.getInstance().getModel(config);
        return faceDetModel;
    }


    private static FaceRecModel faceRecModel = null;
    /**
     * 获取人脸识别模型(带SQLite数据库配置)
     * @return
     */
    public static FaceRecModel getFaceRecModelWithSQLiteConfig() {
        if (faceRecModel != null) {
            return faceRecModel;
        }
        FaceRecConfig config = new FaceRecConfig();
        //高精度模型，速度慢, 追求速度请更换高速模型，具体其他模型参数可以查看文档：http://doc.smartjavaai.cn/face.html
        config.setModelEnum(FaceRecModelEnum.ELASTIC_FACE_MODEL);//人脸检测模型
        config.setModelPath(MODAL_PATH + "elasticface.pt");
        //裁剪人脸：如果图片已经是裁剪过的，则请将此参数设置为false
        config.setCropFace(true);
        //开启人脸对齐：适用于人脸不正的场景，开启将提升人脸特征准确度，关闭可以提升性能
        config.setAlign(true);
        //指定人脸检测模型，可切换人脸检测模型（极速：getFastFaceDetModel，高精度：getProFaceDetModel），具体其他模型参数可以查看文档：http://doc.smartjavaai.cn/face.html
        config.setDetectModel(getProFaceDetModel());
        config.setDevice(device);

        //初始化SQLite数据库
        SQLiteConfig vectorDBConfig = new SQLiteConfig();
        vectorDBConfig.setSimilarityType(SimilarityType.IP);
        config.setVectorDBConfig(vectorDBConfig);
        faceRecModel = FaceRecModelFactory.getInstance().getModel(config);
        return faceRecModel;
    }

    //=================================//=================================//=================================//=================================

    /**
     * 人脸比对1：1（基于图像直接比对）
     * 流程：从输入图像中裁剪分数最高的人脸 → 提取其人脸特征 → 比对两张图片中提取的人脸特征。（接口内自动完成）
     * 注意事项：
     * 1、首次调用接口，可能会较慢。只要不关闭程序，后续调用会明显加快。若每次重启程序，则每次首次调用都将重新加载，仍会较慢。
     * 2、若人脸朝向不正，可开启人脸对齐以提升特征提取准确度。（方法参考自定义配置人脸特征提取）
     * @throws Exception
     */
    public static void featureComparison() {
        try {
            System.out.println(new Date().getSeconds());
            //getHighAccuracyFaceRecModel高精度模型，速度慢, 追求速度请更换高速模型: getHighSpeedFaceRecModel
            FaceRecModel faceRecModel = getFaceRecModelWithSQLiteConfig();
            //基于图像直接比对人脸特征
            R<Float> similarResult = faceRecModel.featureComparison("src/main/resources/iu_1.jpg","src/main/resources/iu_2.jpg");
            if (similarResult.isSuccess()) {
                // 相似度阈值不同模型不同，具体参看文档
                log.info("人脸比对相似度：{}", objectMapper.writeValueAsString(similarResult.getData()));
            } else {
                log.info("人脸比对失败：{}", similarResult.getMessage());
            }
            System.out.println(new Date().getSeconds());
        } catch (Exception e) {
            log.error("人脸特征比对异常: {}", e.getMessage(), e);
        }
    }
    /**
     * 人脸比对1：1（基于特征值比对）
     * 流程：从输入图像中裁剪分数最高的人脸 → 提取其人脸特征 → 比对两张图片中提取的人脸特征。
     * 注意事项：
     * 1、首次调用接口，可能会较慢。只要不关闭程序，后续调用会明显加快。若每次重启程序，则每次首次调用都将重新加载，仍会较慢。
     * 2、若人脸朝向不正，可开启人脸对齐以提升特征提取准确度。（方法参考自定义配置人脸特征提取）
     * @throws Exception
     */
    public static void featureComparison2() {
        try {
            //高精度模型，速度慢, 追求速度请更换高速模型: getHighSpeedFaceRecModel
            FaceRecModel faceRecModel = getFaceRecModelWithSQLiteConfig();
            //特征提取（提取分数最高人脸特征）,适用于单人脸场景
            R<float[]> featureResult1 = faceRecModel.extractTopFaceFeature("src/main/resources/iu_1.jpg");
            if (featureResult1.isSuccess()) {
                log.info("图片1人脸特征提取成功：{}", objectMapper.writeValueAsString(featureResult1.getData()));
            } else {
                log.info("图片1人脸特征提取失败：{}", featureResult1.getMessage());
                return;
            }
            // 特征提取（提取分数最高人脸特征）,适用于单人脸场景
            R<float[]> featureResult2 = faceRecModel.extractTopFaceFeature("src/main/resources/iu_2.jpg");
            if (featureResult2.isSuccess()) {
                log.info("图片2人脸特征提取成功：{}", objectMapper.writeValueAsString(featureResult2.getData()));
            } else {
                log.info("图片2人脸特征提取失败：{}", featureResult2.getMessage());
                return;
            }
            //计算相似度
            float similar = faceRecModel.calculSimilar(featureResult1.getData(), featureResult2.getData());
            log.info("相似度：{}", similar);
        } catch (Exception e) {
            log.error("人脸特征比对异常: {}", e.getMessage(), e);
        }
    }


}
